package com.xuyuan.ai.chat.advisor;


import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Map;

/**
 * TODO 后续阅读了RAG和ChatMemory相关的内容再写案例
 *
 * @author xuyuan
 * <p>
 * @see org.springframework.ai.chat.client.ChatClient.AdvisorSpec
 * @see org.springframework.ai.chat.client.advisor.api.Advisor
 * Advisors API提供了一种灵活而强大的方法来拦截、修改和增强 Spring 应用程序中的 AI 驱动的交互。
 * 使用User Prompt调用 AI 模型时的一个常见模式是使用上下文数据附加或扩充提示。
 * 这些上下文数据可以是不同类型的。常见类型包括：
 * 您自己的数据：这是 AI 模型尚未训练过的数据。即使模型已经看到过类似的数据，附加的上下文数据也会优先生成响应。
 * 对话历史记录：聊天模型的 API 是无状态的。如果您告诉 AI 模型您的姓名，它不会在后续交互中记住它。每次请求都必须发送对话历史记录，以确保在生成响应时考虑到先前的交互。
 * <p>
 * Advisor 添加到链中的顺序至关重要，因为它决定了Advisor的执行顺序。每个Advisor都会以某种方式修改提示或上下文，一个Advisor所做的更改会传递给链中的下一个Advisor。(责任链模式)
 * Spring建议在构建时使用构建器的defaultAdvisors()方法注册。
 * <p> Advisor的核心
 * @see org.springframework.ai.chat.client.advisor.api.CallAroundAdvisor
 * @see org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisor StreamAroundAdvisor和StreamAroundAdvisorChain用于流式场景
 * <p> ChatMemoryAdvisor
 * @see org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor MessageChatMemoryAdvisor 检索内存并作为消息集合添加到提示中，可维护对话历史记录的结构。请注意，并非所有 AI 模型都支持此方法。
 * @see org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor  PromptChatMemoryAdvisor 检索内存并将其添加到提示的系统文本中。
 * @see org.springframework.ai.chat.client.advisor.VectorStoreChatMemoryAdvisor VectorStoreChatMemoryAdvisor 从 VectorStore 检索内存并将其添加到提示的系统文本中，对于高效地从大型数据集中搜索和检索相关信息非常有用。
 * <p> 其他内置的Advisor
 * @see org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor
 * @see org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor QuestionAnswerAdvisor 使用向量存储来提供问答功能，实现 RAG（检索增强生成）模式。
 * @see org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor SimpleLoggerAdvisor记录request和response数据的。这对于调试和监控您的 AI 交互非常有用
 * @see org.springframework.ai.chat.client.advisor.DefaultAroundAdvisorChain
 * @see org.springframework.ai.chat.client.advisor.SafeGuardAdvisor SafeGuardAdvisor防止模型生成有害或不适当的内容
 * <p> 通过CallAroundAdvisor和StreamAroundAdvisor自定义Advisor
 */
@RestController
public class AdvisorController {

    private static final Logger LOGGER = LoggerFactory.getLogger(AdvisorController.class);
    private ChatClient chatClient;

    public AdvisorController(@Qualifier("openAiChatModel") ChatModel openAiChatModel) {
        this.chatClient = ChatClient.builder(openAiChatModel).build();
    }

    /**
     * SimpleLoggerAdvisor记录request和response数据的。
     * 通过配置logging.level.org.springframework.ai.chat.client.advisor设置日志级别DEBUG可以查看
     */
    @GetMapping("/ai/advisor-1")
    public String advisor1(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        return chatClient.prompt()
                .advisors(new SimpleLoggerAdvisor())
                .user("Tell me a joke?")
                .call().content();
    }

    /**
     * 自定义SimpleLoggerAdvisor的请求和响应输出数据
     * SimpleLoggerAdvisor(Function<AdvisedRequest, String> requestToString,Function<ChatResponse, String> responseToString, int order)
     */
    @GetMapping("/ai/advisor-2")
    public String advisor2(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        SimpleLoggerAdvisor customLogger = new SimpleLoggerAdvisor(
                request -> "Custom request: " + request.userText(),
                response -> "Custom response: " + response.getResult(), 0
        );
        return chatClient.prompt()
                .advisors(customLogger)
                .user("Tell me a joke?")
                .call().content();
    }

    /**
     * 自定义Advisor
     */
    static class CustomAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

        @Override
        public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
            LOGGER.info("CustomAdvisor非流式处理前处理");
            AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
            LOGGER.info("CustomAdvisor非流式处理后处理");
            return advisedResponse;
        }

        @Override
        public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
            LOGGER.info("CustomAdvisor流式处理前处理");
            Flux<AdvisedResponse> advisedResponseFlux = chain.nextAroundStream(advisedRequest);
            LOGGER.info("CustomAdvisor流式处理后处理");
            return advisedResponseFlux;
        }

        @Override
        public String getName() {
            return this.getClass().getSimpleName();
        }

        /**
         * 值越小优先级越高
         *
         * @return
         */
        @Override
        public int getOrder() {
            return 0;
        }
    }

    /**
     * Re-Reading Improves Reasoning in Large Language Models（重读）
     * prompt示例：
     * {Input_Query}
     * Read the question again: {Input_Query}
     */
    public class ReReadingAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

        /**
         * 处理输入为Re2
         */
        private AdvisedRequest before(AdvisedRequest advisedRequest) {
            return AdvisedRequest.from(advisedRequest)
                    .userText("""
                            {re2_input_query}
                            Read the question again: {re2_input_query}
                            """)
                    .userParams(Map.of("re2_input_query", advisedRequest.userText()))
                    .build();
        }

        @Override
        public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
            return chain.nextAroundCall(this.before(advisedRequest));
        }

        @Override
        public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
            return chain.nextAroundStream(this.before(advisedRequest));
        }

        @Override
        public int getOrder() {
            return 0;
        }

        @Override
        public String getName() {
            return this.getClass().getSimpleName();
        }
    }
}